package utils

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
	"os"
	"path/filepath"
)

// FileEncryptor 文件加密工具
type FileEncryptor struct {
	key []byte
}

// NewFileEncryptor 创建新的文件加密器
func NewFileEncryptor(key string) (FileEncryptor, error) {
	// 确保密钥长度为16, 24或32字节（对应AES-128, AES-192, AES-256）
	var fixedKey []byte
	switch {
	case len(key) <= 16:
		fixedKey = make([]byte, 16)
		copy(fixedKey, key)
	case len(key) <= 24:
		fixedKey = make([]byte, 24)
		copy(fixedKey, key)
	default:
		fixedKey = make([]byte, 32)
		copy(fixedKey, key)
	}
	return FileEncryptor{key: fixedKey}, nil
}

// EncryptFile 加密文件
func (fe *FileEncryptor) EncryptFile(inputFilePath string) (string, error) {
	// 打开输入文件
	inputFile, err := os.Open(inputFilePath)
	if err != nil {
		return "", err
	}
	defer inputFile.Close()

	// 创建输出文件（添加.enc后缀）
	outputFilePath := inputFilePath + ".enc"
	outputFile, err := os.Create(outputFilePath)
	if err != nil {
		return "", err
	}
	defer outputFile.Close()

	// 创建AES加密块
	block, err := aes.NewCipher(fe.key)
	if err != nil {
		return "", err
	}

	// 创建随机IV（初始化向量）
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}

	// 先写入IV到文件开头
	if _, err := outputFile.Write(iv); err != nil {
		return "", err
	}

	// 创建加密流
	stream := cipher.NewCTR(block, iv)

	// 创建缓冲的加密写入器
	writer := &cipher.StreamWriter{S: stream, W: outputFile}

	// 从输入文件复制到加密写入器
	if _, err := io.Copy(writer, inputFile); err != nil {
		return "", err
	}

	// 删除原文件
	if err := os.Remove(inputFilePath); err != nil {
		return "", err
	}

	return outputFilePath, nil
}

// DecryptFile 解密文件
func (fe *FileEncryptor) DecryptFile(inputFilePath string) (string, error) {
	// 检查文件是否有.enc后缀
	if filepath.Ext(inputFilePath) != ".enc" {
		return "", errors.New("输入文件不是加密文件")
	}

	// 打开输入文件
	inputFile, err := os.Open(inputFilePath)
	if err != nil {
		return "", err
	}
	defer inputFile.Close()

	// 创建输出文件（去掉.enc后缀）
	outputFilePath := inputFilePath[:len(inputFilePath)-4]
	outputFile, err := os.Create(outputFilePath)
	if err != nil {
		return "", err
	}
	defer outputFile.Close()

	// 创建AES加密块
	block, err := aes.NewCipher(fe.key)
	if err != nil {
		return "", err
	}

	// 读取文件开头的IV
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(inputFile, iv); err != nil {
		return "", err
	}

	// 创建解密流
	stream := cipher.NewCTR(block, iv)

	// 创建缓冲的解密读取器
	reader := &cipher.StreamReader{S: stream, R: inputFile}

	// 从解密读取器复制到输出文件
	if _, err := io.Copy(outputFile, reader); err != nil {
		return "", err
	}

	return outputFilePath, nil
}

// GenerateKey 生成随机加密密钥（返回base64编码的字符串）
func GenerateKey() (string, error) {
	key := make([]byte, 32) // AES-256
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(key), nil
}
